import torch
import cv2
import numpy as np
import os

# 设置环境变量指定 torch.hub 路径（若有需要）
os.environ['TORCH_HOME'] = 'D:/machine learning/yolov5-master'
# 加载训练好的 YOLOv5 模型
model = torch.hub.load('ultralytics/yolov5', 'custom', path='runs/train/exp12/weights/best.pt')
# 定义类别名称
classes = ['people', 'bike', 'tree', 'blind_way', 'hard_shoulder', 'waring_sign', 'garbage_can', 'dog', 'road_block', 'trafic_cone', 'car', 'pole', 'cat']

def check_overlap(box1, box2):
    """
    检查两个锚框是否重叠
    """
    x1, y1, x2, y2 = box1
    x3, y3, x4, y4 = box2
    return not (x2 < x3 or x1 > x4 or y2 < y3 or y1 > y4)

def process_frame(frame):
    """
    处理单帧图像
    """
    # 使用 YOLOv5 进行目标检测
    results = model(frame)
    detections = results.pandas().xyxy[0]

    # 提取盲道的锚框
    blind_way_boxes = detections[detections['name'] == 'blind_way'][['xmin', 'ymin', 'xmax', 'ymax']].values

    signal_value = 0
    if len(blind_way_boxes) > 0:
        for blind_way_box in blind_way_boxes:
            for _, obj in detections.iterrows():
                obj_box = [obj['xmin'], obj['ymin'], obj['xmax'], obj['ymax']]
                if check_overlap(blind_way_box, obj_box):
                    if obj['name'] in ['bike', 'car']:
                        signal_value = 2
                        break
                    elif obj['name'] in ['tree', 'hard_shoulder', 'waring_sign', 'garbage_can', 'road_block', 'trafic_cone', 'pole']:
                        signal_value = 1
                        break
                    elif obj['name'] in ['people', 'cat', 'dog']:
                        signal_value = 0
                        break
            if signal_value > 0:
                break
    else:
        signal_value = 0

    # 绘制锚框和标签
    for _, detection in detections.iterrows():
        xmin, ymin, xmax, ymax = int(detection['xmin']), int(detection['ymin']), int(detection['xmax']), int(detection['ymax'])
        label = detection['name']
        cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), (0, 255, 0), 2)  # 绘制绿色锚框，线宽为2
        cv2.putText(frame, label, (xmin, ymin - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)  # 绘制类别标签

    return signal_value

# 示例：读取视频文件并处理每一帧
cap = cv2.VideoCapture('视频3.mp4')

while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break

    signal = process_frame(frame)
    print(f"当前帧的信号值: {signal}")

    # 显示处理后的帧
    cv2.imshow('YOLOv5 Detection', frame)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()